package plugin

import (
	"context"
	"fmt"

	"github.com/louloulin/dataflare/pkg/message"
	"github.com/tetratelabs/wazero/api"
)

// WasmMemory 提供WebAssembly内存操作的辅助功能
// 这个结构体在实际使用wazero库时会有更多的功能
type WasmMemory struct {
	module         api.Module   // wazero的Module对象
	allocateFunc   api.Function // 分配内存函数
	deallocateFunc api.Function // 释放内存函数
}

// NewWasmMemory 创建一个新的WebAssembly内存管理器
func NewWasmMemory(module api.Module, allocateFunc, deallocateFunc api.Function) *WasmMemory {
	return &WasmMemory{
		module:         module,
		allocateFunc:   allocateFunc,
		deallocateFunc: deallocateFunc,
	}
}

// WriteString 将字符串写入WebAssembly内存
// 返回内存指针和字符串长度
func (m *WasmMemory) WriteString(ctx context.Context, s string) (uint32, uint32, error) {
	// 检查allocateFunc是否可用
	if m.allocateFunc == nil {
		return 0, 0, fmt.Errorf("allocate function not available")
	}

	// 将字符串转换为字节数组
	data := []byte(s)
	size := uint32(len(data))

	// 调用WASM中的内存分配函数
	results, err := m.allocateFunc.Call(ctx, uint64(size))
	if err != nil {
		return 0, 0, fmt.Errorf("failed to allocate memory: %w", err)
	}

	if len(results) == 0 {
		return 0, 0, fmt.Errorf("allocate function returned no result")
	}

	// 获取内存指针
	ptr := uint32(results[0])
	if ptr == 0 {
		return 0, 0, fmt.Errorf("memory allocation failed (returned nullptr)")
	}

	// 将数据写入WASM内存
	if !m.module.Memory().Write(ptr, data) {
		return 0, 0, fmt.Errorf("failed to write to memory at address %d", ptr)
	}

	return ptr, size, nil
}

// ReadString 从WebAssembly内存读取字符串
func (m *WasmMemory) ReadString(ctx context.Context, ptr, length uint32) (string, error) {
	// 从WASM内存中读取数据
	data, ok := m.module.Memory().Read(ptr, length)
	if !ok {
		return "", fmt.Errorf("failed to read from memory at address %d, length %d", ptr, length)
	}

	// 将字节数组转换为字符串
	return string(data), nil
}

// FreeMemory 释放WebAssembly内存
func (m *WasmMemory) FreeMemory(ctx context.Context, ptr, length uint32) error {
	// 检查deallocateFunc是否可用
	if m.deallocateFunc == nil {
		return fmt.Errorf("deallocate function not available")
	}

	// 调用WASM中的内存释放函数
	_, err := m.deallocateFunc.Call(ctx, uint64(ptr), uint64(length))
	if err != nil {
		return fmt.Errorf("failed to deallocate memory: %w", err)
	}

	return nil
}

// WriteMessage 将DataMessage写入WebAssembly内存
func (m *WasmMemory) WriteMessage(ctx context.Context, msg message.DataMessage) (uint32, uint32, error) {
	// 序列化消息
	jsonData, err := message.SerializeMessage(msg)
	if err != nil {
		return 0, 0, fmt.Errorf("failed to serialize message: %w", err)
	}

	// 写入内存 - 将 []byte 转换为 string
	return m.WriteString(ctx, string(jsonData))
}

// ReadMessage 从WebAssembly内存读取DataMessage
func (m *WasmMemory) ReadMessage(ctx context.Context, ptr, length uint32) (message.DataMessage, error) {
	// 从内存读取JSON字符串
	jsonStr, err := m.ReadString(ctx, ptr, length)
	if err != nil {
		return nil, fmt.Errorf("failed to read string from memory: %w", err)
	}

	// 反序列化消息 - 将 string 转换为 []byte
	msgInterface, err := message.DeserializeMessage([]byte(jsonStr))
	if err != nil {
		return nil, fmt.Errorf("failed to deserialize message: %w", err)
	}

	// 确保返回的消息实现了 DataMessage 接口
	if msg, ok := msgInterface.(message.DataMessage); ok {
		return msg, nil
	}

	// 如果不是 DataMessage 类型，则创建一个新的 BasicDataMessage 并复制数据
	basicMsg := message.CreateBasicMessage()

	// 尝试将反序列化的结果转换为 map 并设置数据
	if dataMap, ok := msgInterface.(map[string]interface{}); ok {
		// 设置数据字段
		if data, ok := dataMap["data"]; ok {
			basicMsg.SetData(data)
		}

		// 设置元数据
		if metadata, ok := dataMap["metadata"].(map[string]interface{}); ok {
			basicMsg.SetMetadata(metadata)
		}

		// 设置 ID
		if id, ok := dataMap["id"].(string); ok {
			basicMsg.ID = id
		}
	} else {
		// 如果不是 map，则直接设置为数据
		basicMsg.SetData(msgInterface)
	}

	return basicMsg, nil
}

// 以下是实际使用wazero时的实现示例（注释掉，因为当前没有依赖wazero）

/*
func actualWriteString(ctx context.Context, module wazero.Module, allocateFunc wazero.Function, s string) (uint32, uint32, error) {
	// 将字符串转换为字节数组
	data := []byte(s)
	size := uint32(len(data))

	// 调用WASM中的内存分配函数
	results, err := allocateFunc.Call(ctx, uint64(size))
	if err != nil {
		return 0, 0, fmt.Errorf("failed to allocate memory: %w", err)
	}

	// 获取内存指针
	ptr := uint32(results[0])

	// 将数据写入WASM内存
	if !module.Memory().Write(uint32(ptr), data) {
		return 0, 0, fmt.Errorf("failed to write to memory")
	}

	return ptr, size, nil
}

func actualReadString(ctx context.Context, module wazero.Module, ptr, length uint32) (string, error) {
	// 从WASM内存中读取数据
	data, ok := module.Memory().Read(ptr, length)
	if !ok {
		return "", fmt.Errorf("failed to read from memory")
	}

	// 将字节数组转换为字符串
	return string(data), nil
}

func actualFreeMemory(ctx context.Context, deallocateFunc wazero.Function, ptr, length uint32) error {
	_, err := deallocateFunc.Call(ctx, uint64(ptr), uint64(length))
	if err != nil {
		return fmt.Errorf("failed to deallocate memory: %w", err)
	}
	return nil
}
*/

// 以下是为WasmPlugin结构体添加的内存管理方法

// RealProcessImpl 是一个更接近实际wazero实现的处理函数
func (p *WasmPlugin) RealProcessImpl(ctx context.Context, input message.DataMessage) (message.DataMessage, error) {
	// 此函数展示使用wazero时，如何处理WASM内存管理
	if !p.started {
		return nil, fmt.Errorf("plugin not started")
	}

	// 检查process函数是否存在
	if p.processFunc == nil {
		return nil, fmt.Errorf("process function not found in WASM module")
	}

	// 序列化输入消息
	jsonData, err := message.SerializeMessage(input)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize input message: %w", err)
	}

	// 将JSON数据写入WASM内存
	inputPtr, inputLen, err := p.wasmMemory.WriteString(ctx, string(jsonData))
	if err != nil {
		return nil, fmt.Errorf("failed to write data to WASM memory: %w", err)
	}

	// 确保在函数退出时释放内存
	defer p.wasmMemory.FreeMemory(ctx, inputPtr, inputLen)

	fmt.Printf("Process: 调用WASM函数处理数据，传入指针=%d，长度=%d\n", inputPtr, inputLen)

	// 调用WASM中的process函数
	results, err := p.processFunc.Call(ctx, uint64(inputPtr), uint64(inputLen))
	if err != nil {
		return nil, fmt.Errorf("error calling process function: %w", err)
	}

	if len(results) == 0 {
		return nil, fmt.Errorf("process function returned no results")
	}

	// 检查是否返回错误码
	if len(results) == 1 && results[0] < 0 {
		return nil, fmt.Errorf("process function returned error code: %d", results[0])
	}

	var resultPtr, resultLen uint32

	// 处理结果，根据返回值格式可能有两种情况
	if len(results) == 1 {
		// 在单个int64中打包了指针和长度（高32位是长度，低32位是指针）
		packedResult := results[0]
		resultPtr = uint32(packedResult & 0xFFFFFFFF)
		resultLen = uint32(packedResult >> 32)
	} else if len(results) == 2 {
		// 分开返回指针和长度
		resultPtr = uint32(results[0])
		resultLen = uint32(results[1])
	} else {
		return nil, fmt.Errorf("process function returned unexpected result format")
	}

	// 检查结果的有效性
	if resultPtr == 0 || resultLen == 0 {
		return nil, fmt.Errorf("process function returned invalid result pointer or length")
	}

	// 从WASM内存读取JSON结果
	resultJson, err := p.wasmMemory.ReadString(ctx, resultPtr, resultLen)
	if err != nil {
		return nil, fmt.Errorf("failed to read result from WASM memory: %w", err)
	}

	// 释放结果内存
	defer p.wasmMemory.FreeMemory(ctx, resultPtr, resultLen)

	// 反序列化结果
	msgInterface, err := message.DeserializeMessage([]byte(resultJson))
	if err != nil {
		return nil, fmt.Errorf("failed to deserialize result message: %w", err)
	}

	// 确保返回的消息实现了 DataMessage 接口
	if msg, ok := msgInterface.(message.DataMessage); ok {
		return msg, nil
	}

	// 如果不是 DataMessage 类型，则创建一个新的 BasicDataMessage 并复制数据
	basicMsg := message.CreateBasicMessage()

	// 尝试将反序列化的结果转换为 map 并设置数据
	if dataMap, ok := msgInterface.(map[string]interface{}); ok {
		// 设置数据字段
		if data, ok := dataMap["data"]; ok {
			basicMsg.SetData(data)
		}

		// 设置元数据
		if metadata, ok := dataMap["metadata"].(map[string]interface{}); ok {
			basicMsg.SetMetadata(metadata)
		}

		// 设置 ID
		if id, ok := dataMap["id"].(string); ok {
			basicMsg.ID = id
		}
	} else {
		// 如果不是 map，则直接设置为数据
		basicMsg.SetData(msgInterface)
	}

	return basicMsg, nil
}
